from flask import Flask, request, jsonify, render_template, send_from_directory
import os
import subprocess
import tempfile
import json
import logging
import uuid
import shutil
from datetime import datetime
import re
from pathlib import Path
import requests
import torch
import numpy as np
import patchcore.patchcore
import patchcore.common
import patchcore.utils
import patchcore.sampler
from src.new.mvtec_new import SingleImageDataset, ThresholdMethod, BBSNetSaliencyDetector
import psutil
import GPUtil
from functools import lru_cache
import threading
import queue
import time
import cProfile
import pstats
import io
from PIL import Image
import cv2

app = Flask(__name__, static_folder='static', template_folder='templates')

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 配置
UPLOAD_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'static', 'uploads')
RESULTS_FOLDER = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'static', 'results')

# 确保目录存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULTS_FOLDER, exist_ok=True)

# 默认模型路径 - 根据您的环境调整
DEFAULT_MODEL_PATH = "F:\\workspace\\python\\new\\patchcore-inspection-main\\model\\MVTecAD_Results\\IM384_WR50_L2-3_P01_D1024-1024_PS-3_AN-1_S0\\models\\mvtec_blade0"

# BBSNet 模型路径 - 根据您的环境调整
# BBSNET_MODEL_PATH = "model_pt/BBSNet_deploy.pt"

# 设置 src.new.mvtec_new 模块的 BBSNET_MODEL_PATH
# import src.new.mvtec_new
# src.new.mvtec_new.BBSNET_MODEL_PATH = BBSNET_MODEL_PATH
# logger.info(f"设置 BBSNet 模型路径为: {BBSNET_MODEL_PATH}")

# 全局变量存储 PatchCore 模型
patchcore_model = None
device = None
model_lock = threading.Lock()
preprocessing_queue = queue.Queue()
result_cache = {}

# GPU 内存管理配置
GPU_MEMORY_THRESHOLD = 0.8  # GPU 内存使用率阈值
MIN_BATCH_SIZE = 4
MAX_BATCH_SIZE = 8
DEFAULT_BATCH_SIZE = 6

# 缓存配置
CACHE_SIZE = 100
CACHE_TTL = 300  # 缓存有效期（秒）

def time_function(func):
    """简单的函数执行时间装饰器"""
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        logger.info(f"{func.__name__} 执行时间: {end_time - start_time:.2f} 秒")
        return result
    return wrapper

def get_gpu_memory_info():
    """获取 GPU 内存使用情况"""
    try:
        gpus = GPUtil.getGPUs()
        if gpus:
            gpu = gpus[0]  # 使用第一个 GPU
            return {
                'total': gpu.memoryTotal,
                'used': gpu.memoryUsed,
                'free': gpu.memoryFree,
                'utilization': gpu.memoryUtil
            }
        return None
    except Exception as e:
        logger.warning(f"获取 GPU 内存信息失败: {str(e)}")
        return None

def optimize_gpu_memory():
    """优化 GPU 内存使用"""
    try:
        torch.cuda.empty_cache()
        if torch.cuda.is_available():
            torch.cuda.synchronize()
    except Exception as e:
        logger.warning(f"GPU 内存优化失败: {str(e)}")

def calculate_optimal_batch_size():
    """计算最优批处理大小"""
    gpu_info = get_gpu_memory_info()
    if not gpu_info:
        return DEFAULT_BATCH_SIZE

    memory_utilization = gpu_info['utilization']
    if memory_utilization > GPU_MEMORY_THRESHOLD:
        current_batch_size = max(MIN_BATCH_SIZE, DEFAULT_BATCH_SIZE // 2)
    else:
        current_batch_size = min(MAX_BATCH_SIZE, DEFAULT_BATCH_SIZE * 2)

    logger.info(f"当前 GPU 内存使用率: {memory_utilization:.2%}, 批处理大小调整为: {current_batch_size}")
    return current_batch_size

@lru_cache(maxsize=1)
@time_function
def initialize_patchcore():
    """初始化 PatchCore 模型（使用缓存确保只初始化一次）"""
    global patchcore_model, device
    
    with model_lock:
        if patchcore_model is None:
            device = patchcore.utils.set_torch_device([0])
            nn_method = patchcore.common.FaissNN(False, 8)
            patchcore_model = patchcore.patchcore.PatchCore(device)
            patchcore_model.load_from_path(
                load_path=DEFAULT_MODEL_PATH,
                device=device,
                nn_method=nn_method
            )
            logger.info("PatchCore model initialized successfully")
    
    return patchcore_model

@time_function
def preprocess_image(rgb_path, depth_path, resize, imagesize, depth_threshold, threshold_method=ThresholdMethod.OTSU):
    """预处理图像（异步）"""
    try:
        dataset = SingleImageDataset(
            rgb_path=rgb_path,
            depth_path=depth_path,
            classname="unknown",
            resize=resize,
            imagesize=imagesize,
            depth_threshold=depth_threshold,
            threshold_method=threshold_method
        )
        
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=1,
            shuffle=False,
            num_workers=0,
            pin_memory=True
        )
        
        return dataloader
    except Exception as e:
        logger.error(f"图像预处理失败: {str(e)}")
        return None

def get_cached_result(image_hash):
    """获取缓存的结果"""
    if image_hash in result_cache:
        timestamp, result = result_cache[image_hash]
        if time.time() - timestamp < CACHE_TTL:
            return result
        else:
            del result_cache[image_hash]
    return None

def cache_result(image_hash, result):
    """缓存结果"""
    if len(result_cache) >= CACHE_SIZE:
        # 删除最旧的缓存
        oldest_key = min(result_cache.keys(), key=lambda k: result_cache[k][0])
        del result_cache[oldest_key]
    
    result_cache[image_hash] = (time.time(), result)

@time_function
def generate_depth_image(rgb_image_path):
    """调用深度图生成API"""
    url = "http://localhost:5005/generate_depth"
    rgb_image_path = Path(rgb_image_path).as_posix()
    data = {
        "image_path": rgb_image_path
    }

    try:
        response = requests.post(url, json=data)
        response.raise_for_status()

        rgb_filename = os.path.basename(rgb_image_path)
        depth_filename = os.path.splitext(rgb_filename)[0] + '.png'
        depth_path = os.path.join('F:/fan/uploadPath/upload/depth', depth_filename)

        if os.path.exists(depth_path):
            static_depth_dir = os.path.join(app.static_folder, 'depth')
            os.makedirs(static_depth_dir, exist_ok=True)
            static_depth_path = os.path.join(static_depth_dir, depth_filename)
            shutil.copy(depth_path, static_depth_path)
            return depth_path
        else:
            raise Exception("深度图生成失败")

    except Exception as e:
        raise Exception(f"调用深度图生成API失败: {str(e)}")

@app.route('/results/<path:filename>')
def results(filename):
    """提供结果文件访问"""
    return send_from_directory(RESULTS_FOLDER, filename)

@app.route('/health', methods=['GET'])
def health_check():
    """健康检查端点"""
    gpu_info = get_gpu_memory_info()
    return jsonify({
        'status': 'ok',
        'gpu_memory': gpu_info,
        'cache_size': len(result_cache)
    })

@app.route('/detect', methods=['POST'])
@time_function
def detect():
    """处理单张RGB图像的深度图生成和异常检测"""
    try:
        # 检查是否有RGB图像
        if 'rgb_image' not in request.files:
            return jsonify({'error': '未上传RGB图像'}), 400

        rgb_file = request.files['rgb_image']
        if rgb_file.filename == '':
            return jsonify({'error': '未选择RGB图像'}), 400

        # 创建临时目录保存图像
        temp_dir = os.path.join(app.static_folder, 'temp')
        os.makedirs(temp_dir, exist_ok=True)

        # 保存RGB图像
        rgb_filename = os.path.basename(rgb_file.filename)
        rgb_path = os.path.join(temp_dir, rgb_filename)
        rgb_file.save(rgb_path)

        # 计算图像哈希用于缓存
        image_hash = str(hash(rgb_path))
        cached_result = get_cached_result(image_hash)
        if cached_result:
            return jsonify(cached_result)

        # 获取深度图
        depth_dir = 'F:/fan/uploadPath/upload/depth'
        depth_filename = os.path.splitext(rgb_filename)[0] + '.png'
        preexisting_depth_path = os.path.join(depth_dir, depth_filename)

        if os.path.exists(preexisting_depth_path):
            depth_path = preexisting_depth_path
            logger.info(f"使用已有深度图: {depth_path}")
        else:
            try:
                depth_path = generate_depth_image(rgb_path)
            except Exception as e:
                return jsonify({'error': str(e)}), 500

        # 获取参数
        resize = int(request.form.get('resize', '512'))
        imagesize = int(request.form.get('imagesize', '384'))
        depth_threshold = float(request.form.get('depth_threshold', '0.5'))
        threshold_method = request.form.get('threshold_method', 'otsu').upper()
        
        # 确保阈值方法有效
        try:
            threshold_method = ThresholdMethod[threshold_method]
        except KeyError:
            threshold_method = ThresholdMethod.OTSU
            logger.warning(f"无效的阈值方法，使用默认方法 OTSU")

        # 确保模型已初始化
        model = initialize_patchcore()

        # 优化 GPU 内存
        optimize_gpu_memory()

        # 计算最优批处理大小
        batch_size = calculate_optimal_batch_size()

        # 预处理图像
        dataloader = preprocess_image(rgb_path, depth_path, resize, imagesize, depth_threshold, threshold_method)
        if dataloader is None:
            return jsonify({'error': '图像预处理失败'}), 500

        # 进行预测
        predict_start_time = time.time()
        with torch.no_grad():  # 禁用梯度计算
            scores, segmentations, _, _ = model.predict(dataloader)
            anomaly_score = float(scores[0])
        predict_end_time = time.time()
        logger.info(f"模型预测完成，耗时: {predict_end_time - predict_start_time:.2f} 秒")

        # 准备结果
        result = {
            "image_path": rgb_path,
            "anomaly_score": anomaly_score,
            "batch_size": batch_size,
            "threshold_method": threshold_method.value,
            "gpu_memory": get_gpu_memory_info(),
            "status": 'success'
        }

        # 缓存结果
        cache_result(image_hash, result)

        # 保存结果
        results_dir = os.path.join(app.static_folder, 'results', 'temp')
        os.makedirs(results_dir, exist_ok=True)
        
        with open(os.path.join(results_dir, "single_image_result.json"), "w") as f:
            json.dump(result, f, indent=2)

        return jsonify(result)

    except Exception as e:
        logger.exception("处理过程中发生错误")
        return jsonify({'error': str(e)}), 500

@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html')

# @app.route('/generate_saliency', methods=['POST'])
# @time_function
# def generate_saliency():
#     """使用 BBSNet 模型生成显著性图"""
#     try:
#         # 检查是否有RGB图像
#         if 'rgb_image' not in request.files:
#             return jsonify({'error': '未上传RGB图像'}), 400
#
#         rgb_file = request.files['rgb_image']
#         if rgb_file.filename == '':
#             return jsonify({'error': '未选择RGB图像'}), 400
#
#         # 检查是否有深度图
#         if 'depth_image' not in request.files:
#             return jsonify({'error': '未上传深度图'})
#
#         depth_file = request.files['depth_image']
#         if depth_file.filename == '':
#             return jsonify({'error': '未选择深度图'})
#
#         # 创建临时目录保存图像
#         temp_dir = os.path.join(app.static_folder, 'temp')
#         os.makedirs(temp_dir, exist_ok=True)
#
#         # 保存图像
#         rgb_filename = os.path.basename(rgb_file.filename)
#         rgb_path = os.path.join(temp_dir, rgb_filename)
#         rgb_file.save(rgb_path)
#
#         depth_filename = os.path.basename(depth_file.filename)
#         depth_path = os.path.join(temp_dir, depth_filename)
#         depth_file.save(depth_path)
#
#         # 检查图像是否存在
#         if not os.path.exists(rgb_path) or not os.path.exists(depth_path):
#             return jsonify({'error': '保存图像失败'}), 500
#
#         # 初始化 BBSNet 检测器
#         detector = BBSNetSaliencyDetector()
#
#         # 加载图像
#         rgb_image = Image.open(rgb_path).convert('RGB')
#         depth_image = Image.open(depth_path).convert('L')
#
#         # 生成显著性图
#         saliency_map = detector.detect(rgb_image, depth_image)
#
#         # 保存显著性图
#         saliency_filename = f"saliency_{os.path.splitext(rgb_filename)[0]}.png"
#         saliency_path = os.path.join(temp_dir, saliency_filename)
#
#         # 转换为 uint8 并保存
#         saliency_uint8 = (saliency_map * 255).astype(np.uint8)
#         cv2.imwrite(saliency_path, saliency_uint8)
#
#         # 如果需要二值化
#         threshold = float(request.form.get('threshold', '0.5'))
#         binary_filename = f"binary_{os.path.splitext(rgb_filename)[0]}.png"
#         binary_path = os.path.join(temp_dir, binary_filename)
#
#         _, binary_mask = cv2.threshold(
#             saliency_uint8,
#             int(threshold * 255),
#             255,
#             cv2.THRESH_BINARY
#         )
#         cv2.imwrite(binary_path, binary_mask)
#
#         # 返回结果
#         return jsonify({
#             'status': 'success',
#             'saliency_path': f"/static/temp/{saliency_filename}",
#             'binary_path': f"/static/temp/{binary_filename}"
#         })
#
#     except Exception as e:
#         logger.exception("生成显著性图时出错")
#         return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    # 初始化模型
    # 先初始化 BBSNet 检测器
    # logger.info("初始化 BBSNet 检测器...")
    # from src.new.mvtec_new import BBSNetSaliencyDetector
    # bbsnet = BBSNetSaliencyDetector()
    # logger.info("BBSNet 检测器初始化完成")
    
    # 初始化 PatchCore 模型
    initialize_patchcore()
    
    # 启动 Flask 应用
    app.run(host='0.0.0.0', port=5000, debug=True)